#include "kernels_head.cuh"

__global__ void init_cluster_centers(float *d_anchors, const int num_of_clusters, const int *d_data_row_ptr, const int *d_data_cols_ind, const float *d_data_vals, const int n, const int d){

	int tid = threadIdx.x + blockDim.x * blockIdx.x;
	if (tid < num_of_clusters){
		int start_idx = d_data_row_ptr[tid];
		int end_idx = d_data_row_ptr[tid + 1];
		for (int i = start_idx; i < end_idx; i++){
			d_anchors[d_data_cols_ind[i] + tid*d] = d_data_vals[i];
		}
	}
}

__global__ void get_euclid_square_dist_kernel(float *d_euclid_squared_dist_2, const int num_of_clusters, const int n, const float *d_anchors, const float *d_vals, const int *d_row_ptr, const int *d_cols_ind, const int d){
	int row = threadIdx.y + blockDim.y * blockIdx.y;
	int col = threadIdx.x + blockDim.x * blockIdx.x;
	if (row < num_of_clusters && col < n){
		float euclid_square_dist = 0.0f;
		for (int i = 0; i < d; i++){
			euclid_square_dist += d_anchors[i + row*d] * d_anchors[i + row*d];
		}
		int start_idx = d_row_ptr[col];
		int end_idx = d_row_ptr[col + 1];
		for (int i = start_idx; i < end_idx; i++){
			euclid_square_dist += d_vals[i] * d_vals[i] - 2 * d_vals[i] * d_anchors[d_cols_ind[i] + row*d];
		}
		d_euclid_squared_dist_2[row + col*num_of_clusters] = euclid_square_dist;
	}
}
__global__ void find_nearest_cluster(const float *d_euclid_squared_dist_2, int *d_membership, const int n, const int num_of_clusters){
	int tid = threadIdx.x + blockDim.x * blockIdx.x;
	if (tid < n){
		float min_dist, dist;
		int index;

		index = 0;
		min_dist = d_euclid_squared_dist_2[0 + num_of_clusters*tid];
		for (int i = 1; i < num_of_clusters; i++){
			dist = d_euclid_squared_dist_2[i + num_of_clusters*tid];
			if (dist < min_dist){
				min_dist = dist;
				index = i;
			}
		}

		d_membership[tid] = index;
	}
}

__global__ void update_cluster(float *d_anchors, const int *d_data_row_ptr, const int *d_data_cols_ind, const float *d_data_vals, const int n, const int num_of_clusters, const int d, const int *d_membership){

	int row = threadIdx.y + blockDim.y * blockIdx.y;
	int col = threadIdx.x + blockDim.x * blockIdx.x;

	if (row < d&&col < num_of_clusters){
		float sum = 0.0f;
		int cnt = 0;
		int start_idx, end_idx;
		for (int i = 0; i < n; i++){
			if (d_membership[i] == col){
				cnt++;
				start_idx = d_data_row_ptr[i];
				end_idx = d_data_row_ptr[i + 1];
				for (int j = start_idx; j < end_idx; j++){
					if (d_data_cols_ind[j] == row){
						sum += d_data_vals[j];
						break;
					}
				}
			}
			if (cnt>0)
				d_anchors[row + d*col] = sum / cnt;
		}
	}
}